#!/bin/bash

# =============================================================================
# 文件传输脚本
# 作者: ShineZhang
# 日期: 2024年12月
# 描述: 该脚本用于通过 SSH 和 SCP 上传和下载文件。
#        通过指定本地文件路径和远程主机路径，可以实现文件的传输。
#        支持通过参数选择上传或下载文件。
# =============================================================================

# 配置颜色
RESET_COLOR="\033[0m"
GREEN="\033[32m"
YELLOW="\033[33m"
RED="\033[31m"
BLUE="\033[34m"
BOLD="\033[1m"
UNDERLINE="\033[4m"
CYAN="\033[36m"

# 分隔符
LINE="========================================================"

# 使用帮助信息
function show_help() {
  echo -e "${BOLD}${GREEN}${LINE}${RESET_COLOR}"
  echo -e "${BOLD}${GREEN} 文件传输脚本 ${RESET_COLOR}"
  echo -e "${BOLD}${GREEN}${LINE}${RESET_COLOR}"
  echo -e "用法: $0 [本地文件路径] [选项]"
  echo -e "${BOLD}${YELLOW}选项:${RESET_COLOR}"
  echo -e "  -r              从远程主机下载文件到本地"
  echo -e "  -h, --help      显示此帮助信息"
  echo -e ""
  echo -e "示例:"
  echo -e "  上传文件: $0 /path/to/local/file"
  echo -e "  下载文件: $0 /path/to/local/destination -r"
  echo -e "${BOLD}${GREEN}${LINE}${RESET_COLOR}"
  exit 0
}

# 检查是否提供了文件路径作为输入参数
if [ -z "$1" ]; then
  echo -e "${RED}错误: 请提供本地文件路径作为参数。${RESET_COLOR}"
  show_help
fi

# 本地文件路径从输入参数获取
LOCAL_FILE="$1"

# 远程主机信息（使用前修改）
REMOTE_USER="xxx"              # 远程主机登录用户名
REMOTE_HOST="42.193.219.111"    # 远程主机IP
REMOTE_PATH="/home/joke/scp"    # 远程主机传输路径
SSH_PORT="6000"                 # 远程主机SSH自定义端口，默认22
REMOTE_PASSWORD="xxx"             # 远程主机登录密码

# 打印远程主机信息
function print_remote_info() {
  echo -e "${CYAN}${BOLD}${LINE}${RESET_COLOR}"
  echo -e "${BOLD}${BLUE} 远程主机信息 ${RESET_COLOR}"
  echo -e "${CYAN}${BOLD}${LINE}${RESET_COLOR}"
  
  # 使用 printf 进行对齐，避免显示转义字符
  printf "%-15s: %s\n" "$(echo -e ${UNDERLINE}用 户 名${RESET_COLOR})" "$REMOTE_USER"
  printf "%-15s: %s\n" "$(echo -e ${UNDERLINE}主机地址${RESET_COLOR})" "$REMOTE_HOST"
  printf "%-15s: %s\n" "$(echo -e ${UNDERLINE}远程路径${RESET_COLOR})" "$REMOTE_PATH"
  printf "%-15s: %s\n" "$(echo -e ${UNDERLINE}SSH 端口${RESET_COLOR})" "$SSH_PORT"
  
  echo -e "${CYAN}${BOLD}${LINE}${RESET_COLOR}"
  echo -e ""
}

# 打印远程主机信息
print_remote_info

# 检查并安装 sshpass 和 rsync
if ! command -v sshpass &> /dev/null; then
  echo -e "${YELLOW}警告: sshpass 未安装，正在安装...${RESET_COLOR}"
  sudo apt-get install -y sshpass
fi

if ! command -v rsync &> /dev/null; then
  echo -e "${YELLOW}警告: rsync 未安装，正在安装...${RESET_COLOR}"
  sudo apt-get install -y rsync
fi

# 获取文件名
FILENAME=$(basename "$LOCAL_FILE")

# 判断第二个参数是否为 -r（下载）
if [ "$2" == "-r" ]; then
  # 从远程路径下载文件
  echo -e "${CYAN}正在从远程主机 ${BOLD}$REMOTE_USER@$REMOTE_HOST${RESET_COLOR} 下载文件: ${BOLD}$REMOTE_PATH/$FILENAME${RESET_COLOR} 到本地 ${BOLD}$LOCAL_FILE${RESET_COLOR} ..."
  
  # 使用 sshpass 自动填充密码，并使用 rsync 下载文件
  sshpass -p "$REMOTE_PASSWORD" rsync -avz --progress -e "ssh -p $SSH_PORT" "$REMOTE_USER@$REMOTE_HOST:$REMOTE_PATH/$FILENAME" "$LOCAL_FILE"

  # 检查下载是否成功
  if [ $? -eq 0 ]; then
    echo -e "${GREEN}文件下载成功！${RESET_COLOR}"
  else
    echo -e "${RED}文件下载失败！${RESET_COLOR}"
  fi
else
  # 执行文件上传
  echo -e "${CYAN}正在上传文件 ${BOLD}$LOCAL_FILE${RESET_COLOR} 到远程主机 ${BOLD}$REMOTE_USER@$REMOTE_HOST${RESET_COLOR} 路径: ${BOLD}$REMOTE_PATH/$FILENAME${RESET_COLOR} ..."
  
  # 使用 sshpass 自动填充密码，并使用 rsync 上传文件
  sshpass -p "$REMOTE_PASSWORD" rsync -avz --progress -e "ssh -p $SSH_PORT" "$LOCAL_FILE" "$REMOTE_USER@$REMOTE_HOST:$REMOTE_PATH/$FILENAME"

  # 检查上传是否成功
  if [ $? -eq 0 ]; then
    echo -e "${GREEN}文件上传成功！${RESET_COLOR}"
  else
    echo -e "${RED}文件上传失败！${RESET_COLOR}"
  fi
fi
